# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

import numpy as np
from scipy.fftpack import fft
import matplotlib.pyplot as plt
from scipy import signal


class Dsp:

    def FFT (self,data,Fs):
        L = len (data)                        # 信号长度
        N =int(np.power(2,np.ceil(np.log2(L))))    # 下一个最近二次幂
        FFT_y1 = np.abs(fft(data,N))/L*2      # N点FFT 变化,但处于信号长度
        Fre = np.arange(int(N/2))*Fs/N        # 频率坐标
        FFT_y1 = FFT_y1[range(int(N/2))]      # 取一半
        plt.figure
        plt.title('FFT Magnitude')
        plt.ylabel('Amplitude')
        plt.xlabel('Frequency [Hz]')
        plt.plot(Fre,FFT_y1)
        plt.grid()
        plt.show()
        return Fre, FFT_y1
    
    def STFT(self,y,Fs):
        f,t,Zxx=signal.stft(y,Fs)
        plt.pcolormesh(t, f, np.abs(Zxx), vmin = 0, vmax = 3)
        plt.title('STFT Magnitude')
        plt.ylabel('Frequency [Hz]')
        plt.xlabel('Time [sec]')
        plt.show()
        
    def LOWPASS(self,data,Fs,cutoff_fre):
        wn=2*cutoff_fre/Fs
        b, a = signal.butter(8, wn, 'lowpass')   #配置滤波器 8 表示滤波器的阶数
        filtedData = signal.filtfilt(b, a, data)  #data为要过滤的信号
        return filtedData
        
    def HIGHPASS(self,data,Fs,cutoff_fre):
        wn=2*cutoff_fre/Fs
        b, a = signal.butter(8, wn, 'highpass')   #配置滤波器 8 表示滤波器的阶数
        filtedData = signal.filtfilt(b, a, data)  #data为要过滤的信号
        return filtedData
    
    def BANDPASS(self,data,Fs,cutoff1,cutoff2):
        wn1=2*cutoff1/Fs
        wn2=2*cutoff2/Fs
        b, a = signal.butter(8, [wn1,wn2], 'bandpass')   #配置滤波器 8 表示滤波器的阶数
        filtedData = signal.filtfilt(b, a, data)  #data为要过滤的信号
        return filtedData
    
    def BANDSTOP(self,data,Fs,cutoff1,cutoff2):
        wn1=2*cutoff1/Fs
        wn2=2*cutoff2/Fs
        b, a = signal.butter(8, [wn1,wn2], 'bandstop')   #配置滤波器 8 表示滤波器的阶数
        filtedData = signal.filtfilt(b, a, data)  #data为要过滤的信号
        return filtedData


    
    
    
    














